# data_fetcher.py
"""
负责从外部数据源获取原始文档。
"""
import logging
from typing import List
from llama_index.core.schema import Document
from llama_index.tools.tavily_research import TavilyToolSpec

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def fetch_documents_from_tavily(api_key: str, query: str, max_results: int) -> List[Document]:
    """
    使用 Tavily API 搜索并获取相关文档。

    Args:
        api_key (str): Tavily 的 API 密钥。
        query (str): 搜索查询语句。
        max_results (int): 最大返回结果数量。

    Returns:
        List[Document]: LlamaIndex 文档对象列表。
    """
    logging.info(f"开始使用 Tavily 搜索，查询: '{query}'...")
    try:
        tavily_tool = TavilyToolSpec(api_key=api_key)
        documents = tavily_tool.search(query, max_results=max_results)
        logging.info(f"成功从 Tavily 获取 {len(documents)} 篇文档。")
        return documents
    except Exception as e:
        logging.error(f"从 Tavily 获取数据时发生错误: {e}")
        # 根据需求可以选择抛出异常或返回空列表
        raise